import multiprocessing
import glob
import itertools
import sys
import scipy.stats
import scipy.misc
import numpy
import os
import operator
from potpour import Worker






def stack(D):
    """
    from list of bases at a site D,
    returns an ordered list of counts of bases
    """
    L = len(D)
    counts = []
    for i in range(len(D[0])):
        R=Y=S=W=K=M=0
        for nseq in range(L):
            R += D[nseq][i].count("R")
            Y += D[nseq][i].count("Y")
            S += D[nseq][i].count("S")
            W += D[nseq][i].count("W")
            K += D[nseq][i].count("K")
            M += D[nseq][i].count("M")
        counts.append( [R,Y,S,W,K,M] )
    return counts


def countpolys(seqs):
    t = [tuple(seq) for seq in seqs]
    return max([sum(i) for i in stack(t)])


def polyaccept(maxpoly, MAXpoly, ln):
    passed = 1
    if 'p' in str(MAXpoly):
        if maxpoly > ln*float(MAXpoly.replace('p','')):
            passed = 0
    else:
        MAXpoly = int(MAXpoly)
        if maxpoly > MAXpoly:
            passed = 0
    if passed:
        return 1
    

def ffmin(x):
    d = []
    for i,j in enumerate(x):
        if j not in "-":
            d.append(i)
    return min(d)

def ffmax(x):
    d = []
    for i,j in enumerate(x):
        if j not in "-":
            d.append(i)
    return max(d)


def alignfast(names,seqs,muscle):
    ST = "\n".join('>'+i+'\n'+j[0] for i,j in zip(names,seqs))
    cmd = "echo '"+ST+"' | "+muscle+" -quiet -in -"
    fin, fout = os.popen4(cmd)
    return fout.read()



def sortalign(stringnames):
    G = stringnames.split("\n>")
    GG = [i.split("\n")[0].replace(">","")+"\n"+"".join(i.split('\n')[1:]) for i in G]
    aligned = [i.split("\n") for i in GG]
    nn = [">"+i[0] for i in aligned]
    seqs = [i[1] for i in aligned]
    return nn,seqs


def FUNC(infile, minspecies, MAXpoly, maxSNP, muscle, minlen):
    f = open(infile)
    aout = open(infile+".cons",'w')
    locus = paralog = g4 = dups = 0
    saved = {}
    k = itertools.izip(*[iter(f)]*2)
    while 1:
        D = P = S = notes = ""
        locus += 1
        try: d = k.next()
        except StopIteration : break
        names = []
        tnames = []
        seqs = []
        while d[0] != "//\n":
            "record names and seqs, remove # at end"
            "record the name into locus name. "
            tnames.append("_".join(d[0].split(">")[1].split("_")[:-2]))
            names.append("".join(d[0].split(">")[1])) #.split("_")[:-2]))
            seqs.append([d[1].strip()])
            d = k.next()
        if len(tnames) != len(set(tnames)):      ## no grouping un-clustered copies from same taxon
            dups += 1                          ## record duplicates in loci
            D = '%D'
         if len(names) > (minspecies-1):
            g4 += 1
            stringnames = alignfast(names,seqs,muscle)
            nn, seqs = sortalign(stringnames)
            maxpoly = countpolys(seqs)                              ## count polys on aligned seqs
            if not D:
                if not polyaccept(maxpoly, MAXpoly, len(nn)):
                    P = '%P'
            bases = []
            for i in range(len(seqs[0])):      ## create list of bases at each site
                site = [s[i] for s in seqs]
                bases.append(site)
            basenumber = 0
            consensus = ""
            snpsite = [" " for i in range(len(seqs[0]))]
            for site in bases:
                reals = [i for i in site if i not in list("N-")]
                if len(set(reals)) > 1:                                ## if site is variable
                    if sorted([reals.count(i) for i in set(reals)],reverse=True)[1] > 1: # not autapomorphy
                        snpsite[basenumber] = "*"                      ## mark PIS for outfile .align
                        consensus += "N"
                    else:                                              ## if autapormorphy
                        snpsite[basenumber] = '-'
                        consensus += max(set(reals), key=reals.count)
                elif len(set(reals)) == 1:
                    consensus += reals[0]
                basenumber += 1
            zz = zip(nn,seqs)
            pp = map(ffmin,seqs)
            FM = max(pp)
            pp = map(ffmax,seqs)
            SM = min(pp)
            zz.sort()
            if not (D or P):
                if any([dash.strip("-").lstrip("-").count("-") > 10 for dash in seqs]):
                    S = "%S"     ## bad alignment
                if (snpsite[FM:SM+1].count('*') > maxSNP):
                    S = "%S"
                if (snpsite[FM:SM+1].count('-') > ((SM+1)-FM)/4.0):
                    S = "%S"
                if ((SM+1)-FM) < minlen:
                    S = "%S"
            if len(D+P+S) == 0:
                for x,y in zz:                               ## write to .aligned file
                    if zz[0][0] not in saved:
                        saved[zz[0][0]] = [x,y[FM:SM+1]]
                    else:
                        saved[zz[0][0]].append([x,y[FM:SM+1]])
                    print x[0:14]+" "*(15-len(x[0:14]))+y[FM:SM+1]
                print '//'+D+P+S+' '*(13-len(D+P+S))+"".join(snpsite[FM:SM+1])+"|"+notes
                #print >>aout, zz[0][0][0:14]+" "*(15-len(x[0:14]))+consensus+'\n'
                print >>aout, zz[0][0]+"\n"+consensus[FM:SM]

            # else:
            #     for x,y in zz:                               ## write to .aligned file
            #         print >>nout, x[0:14]+" "*(15-len(x[0:14]))+y
            #     print >>nout, '//'+D+P+S+' '*(13-len(D+P+S))+"".join(snpsite)+"|"
    aout.close()
    print '.'
    #print locus
    #return locus #,aout,nout


infile = str(sys.argv[1])
minspecies = 1
MAXpoly = 3
maxSNP = 10
muscle = 'muscle'
minlen = 50

FUNC(infile, minspecies, MAXpoly, maxSNP, muscle, minlen)

